import hashlib
import math
import numpy
import os
import os.path as osp
import cPickle as pickle
import struct
import subprocess
import sys
import wave

from optparse import OptionParser

import mutagen.easyid3
import mutagen.id3
import mutagen.mp3

import tables

from mp3row import MP3Row

def generate():
    op = OptionParser("mp3same gen [opts] file1.mp3 | dir1 [file2.mp3 | dir2 [...]]")
    op.add_option("-f", dest="fresh", action="store_true", default=False,
                  help="Force overwrite of output file")
    op.add_option('-o', '--output', dest='output', default='mp3same.h5',
                  help='Redirect output to PATH', metavar='PATH')
    op.add_option('-d', '--dir-search', action='store_true', dest='search',
                  default=False, 
                  help='Treat arguments as directories to search')
    op.add_option("-s", dest="samples_sec", type='int', default=50,
                  help="Use N samples per second", metavar='N')

    (opts,args) = op.parse_args()

    if len(args) < 2:
        sys.stderr.write("`mp3same.py gen' requires at least one argument.\n")
        sys.stderr.write("See `mp3same.py gen -h' for more.\n")
        sys.stderr.flush()
        return 1

    pkl_fname = '.'.join(opts.output.rsplit('.')[:-1] + ['pkl'])

    if osp.exists(opts.output):
        if opts.fresh:
            os.remove(opts.output)
            try:
                os.remove(pkl_fname)
            except OSError:
                pass


            trash = list()
            md5dict = dict()
            outfp = tables.openFile(opts.output, mode='w', 
                                    title='mp3stat database file')
            arrays_group = outfp.createGroup(outfp.root, 
                                             'data', 'Condensed WAV data')
            hdrnt = outfp.createTable(outfp.root, 'hdrnt', MP3Row,
                                      "mp3stat file info table")

        else:
            outfp = tables.openFile(opts.output, mode='a')
            hdrnt = outfp.root.hdrnt
            arrays_group = outfp.root.data
            pickfp = open(pkl_fname, 'r')
        
            (md5dict, trash) = pickle.load(pickfp)

            pickfp.close()
    else:
        trash = list()
        md5dict = dict()
        outfp = tables.openFile(opts.output, mode='w', 
                                title='mp3stat database file')
        arrays_group = outfp.createGroup(outfp.root, 
                                         'data', 'Condensed WAV data')
        hdrnt = outfp.createTable(outfp.root, 'hdrnt', MP3Row,
                                  "mp3stat file info table")

    mp3hdr = hdrnt.row
 
    if opts.search:
        file_list = list()
        for path in args[1:]:
            for root, dirs, files in os.walk(path):
                for fname in files:
                    if fname[-4:].lower() == '.mp3':
                        file_list.append(osp.abspath(osp.join(root, fname)))
    else:
        file_list = [osp.abspath(fname) for fname in args[1:] if fname[-4:].lower() == '.mp3']
    file_list.sort()

    files_done = 0
    file_list.sort()

    already_done = set(outfp.root.hdrnt[:]['fname'])

    file_list = [fname for fname in file_list if fname not in already_done]

    try:
        for fname in file_list:
            print files_done, 'of', len(file_list), fname

            md5fp = open(fname, 'rb')
            md5sum = hashlib.md5(md5fp.read()).hexdigest()            

            if md5sum not in md5dict:
                md5dict[md5sum] = [osp.abspath(fname)]
            else:
                md5dict[md5sum].append(osp.abspath(fname))
                files_done += 1
                continue

            mp3hdr['mp3md5'] = md5sum
            md5fp.close()

            mp3hdr['fname']      = fname
            mp3hdr['mp3size']    = os.stat(fname).st_size
            
            try:
                id3tags = mutagen.easyid3.EasyID3(fname)
            except mutagen.id3.ID3NoHeaderError:
                mp3hdr['artist'] = 'unknown'
                mp3hdr['album']  = 'unknown'
                mp3hdr['title']  = 'unknown'
                mp3hdr['track']  = 0
            else:        
                try:
                    mp3hdr['artist'] = id3tags['artist'][0].encode('ascii', 'ignore')
                except KeyError:
                    mp3hdr['artist'] = 'unknown'
                try:
                    mp3hdr['album']   = id3tags['album'][0].encode('ascii', 'ignore')
                except KeyError:
                    mp3hdr['album'] = 'unknown'
                try:
                    mp3hdr['track'] = id3tags['tracknumber'][0].encode('ascii', 'ignore')
                except KeyError:
                    mp3hdr['track'] = 0
                try:
                    mp3hdr['title'] = id3tags['title'][0].encode('ascii', 'ignore')
                except KeyError:
                    mp3hdr['title'] = 'unknown'
                                
            try:
                mp3f = mutagen.mp3.MP3(fname)
            except mutagen.mp3.HeaderNotFoundError:
                trash.append(mp3hdr['fname'])
                files_done += 1
                continue

            mp3hdr['length'] = mp3f.info.length
            mp3hdr['bitrate'] = mp3f.info.bitrate


            wave_out = osp.abspath(fname)[:-4] + ".wav"
        
            lame_args = ['lame', '--decode', osp.abspath(fname), wave_out]

            lame_pipe = subprocess.Popen(lame_args, stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE)

            lame_out, lame_err = lame_pipe.communicate()

            if lame_pipe.returncode != 0:
                print "Error encountered while decoding `" + fname + "':"
                print lame_err 
                print "Skipping..."
                files_done += 1
                continue

            try:
                wavfp = wave.open(wave_out)
            except IOError:
                print "Error opening `" + wave_out + "'"
                print "Skipping..."
                try:
                    os.remove(wave_out)
                except:
                    pass
            
                files_done += 1
                continue

            mp3hdr['nframes'] = wavfp.getnframes()

            if mp3hdr['nframes'] < 500:
                os.remove(wave_out)
                trash.append(mp3hdr['fname'])
                files_done += 1
                continue


            nchannels = wavfp.getnchannels()

            if nchannels != 2:
                os.remove(wave_out)
                files_done += 1
                continue
        

            mp3hdr['nchannels'] = nchannels
            nframes_sec = float(wavfp.getframerate())
            nframes_sample = int(math.floor(nframes_sec / float(opts.samples_sec)))  
            wavstruct = struct.Struct(str(nframes_sample*nchannels) + 'h')

            wavdata = numpy.zeros((mp3hdr['nframes'] / nframes_sample + 1.,))
        
            ind = 0
        
            while True:            
                samplestr = wavfp.readframes(nframes_sample)
                if len(samplestr) != wavstruct.size:
                    if len(samplestr):
                        data = struct.unpack(str(len(samplestr)/2) + 'h', samplestr)
                        wavdata[ind] = float(sum(data)) / float(len(samplestr))
                        break
                    else:
                        break

                sample_data = wavstruct.unpack(samplestr)
                wavdata[ind] = float(sum(sample_data)) / float(nframes_sample)
                ind += 1

            arrname = "arr" + md5sum
            mp3hdr['arrname'] = arrname
            outfp.createArray(arrays_group, arrname, wavdata,
                          arrname + ' samples')

            mp3hdr.append()
            hdrnt.flush()
       
            os.remove(wave_out)
            files_done += 1

    except KeyboardInterrupt:
        print "Closing output file and exiting..."
        print "You may have to clean up stray WAV files yourself."

    finally:
        outfp.close()
        
        pickfp = open(pkl_fname, 'wb')
        pickle.dump((md5dict, trash), pickfp)
        
        outfp.close()
